import rclpy
from rclpy.node import Node
import numpy as np
from sensor_msgs.msg import PointCloud2, PointField
from std_msgs.msg import Header
import struct
import math

# 带颜色的点云字段定义
COLORED_FIELDS = [
    PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
    PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
    PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1),
    PointField(name='r', offset=12, datatype=PointField.UINT8, count=1),
    PointField(name='g', offset=13, datatype=PointField.UINT8, count=1),
    PointField(name='b', offset=14, datatype=PointField.UINT8, count=1),
]
COLORED_POINT_STEP = 15  # 每个点的字节数


class CablePointCloudSimulator(Node):
    def __init__(self):
        super().__init__('cable_pointcloud_simulator')
        self.publisher_ = self.create_publisher(PointCloud2, 'simulated_pointcloud', 10)
        self.timer = self.create_timer(0.1, self.timer_callback)  # 10Hz发布
        
        # 激光雷达配置（向下朝向）
        self.lidar_direction = "downward"  # 向下扫描
        
        # 钢缆参数（2根钢缆）
        self.cable_length = 10.0  # 钢缆长度（米）
        self.cable_diameter = 0.02  # 钢缆直径（20mm）
        self.swing_angle_theta = 10  # 与竖直方向夹角（度）
        self.cable_spacing = 1.0       # 双钢缆间距（米）
        self.noise_level = 0.1       # 点云噪声幅度
        
        # 点云密度参数（总点数≈1000）
        self.cable_points_per_meter = 10   # 每米钢缆的主线点数
        self.radial_points = 3             # 每个主线点的径向采样点
        self.cargo_points = 2000            # 货物点数

    def generate_single_cable(self, x_offset):
        """生成单根带直径的钢缆点云（带X方向偏移，左右分布）"""
        # 角度转换为弧度
        theta = math.radians(self.swing_angle_theta)
        
        # 激光雷达向下时，钢缆方向向量（z轴向下为正）
        dir_x = 0.0
        dir_y = 0.0
        dir_z = 1.0  # 竖直向下
        
        # 加入摆动偏移（轻微摆动）
        swing_phi = np.random.uniform(-theta/2, theta/2)  # 水平摆动角度
        dir_x = math.sin(theta) * math.cos(swing_phi)
        dir_y = math.sin(theta) * math.sin(swing_phi)
        
        # 计算钢缆主线点数
        num_line_points = int(self.cable_length * self.cable_points_per_meter)
        t = np.linspace(0, self.cable_length, num_line_points)
        
        # 生成主线点（带X方向偏移，左右分布）
        line_x = x_offset + dir_x * t
        line_y = 0.0 + dir_y * t  # Y方向居中
        line_z = 0.0 + dir_z * t  # 从原点（雷达位置）向下延伸
        
        # 生成垂直于钢缆方向的两个正交向量（用于模拟直径）
        # 向量1：水平垂直于钢缆方向
        ortho1_x = -dir_y if abs(dir_y) > 0.01 else 1.0
        ortho1_y = dir_x if abs(dir_x) > 0.01 else 0.0
        ortho1_z = 0.0
        
        # 向量2：垂直于钢缆和向量1的方向
        ortho2_x = dir_y * ortho1_z - dir_z * ortho1_y
        ortho2_y = dir_z * ortho1_x - dir_x * ortho1_z
        ortho2_z = dir_x * ortho1_y - dir_y * ortho1_x
        
        # 归一化正交向量
        ortho1_norm = math.sqrt(ortho1_x**2 + ortho1_y**2 + ortho1_z**2)
        ortho2_norm = math.sqrt(ortho2_x**2 + ortho2_y**2 + ortho2_z**2)
        
        ortho1_x /= ortho1_norm
        ortho1_y /= ortho1_norm
        ortho1_z /= ortho1_norm
        
        ortho2_x /= ortho2_norm
        ortho2_y /= ortho2_norm
        ortho2_z /= ortho2_norm
        
        # 生成径向点（模拟直径）
        cable_points = []
        radius = self.cable_diameter / 2  # 半径
        
        for i in range(num_line_points):
            # 主线点坐标
            x0, y0, z0 = line_x[i], line_y[i], line_z[i]
            
            # 在径向生成点（圆形分布）
            for r in range(self.radial_points):
                angle = 2 * math.pi * r / self.radial_points
                # 计算径向偏移
                dx = radius * (math.cos(angle) * ortho1_x + math.sin(angle) * ortho2_x)
                dy = radius * (math.cos(angle) * ortho1_y + math.sin(angle) * ortho2_y)
                dz = radius * (math.cos(angle) * ortho1_z + math.sin(angle) * ortho2_z)
                
                # 添加带噪声的点
                cable_points.append([
                    x0 + dx + np.random.normal(0, self.noise_level),
                    y0 + dy + np.random.normal(0, self.noise_level),
                    z0 + dz + np.random.normal(0, self.noise_level/2)
                ])
        
        return np.array(cable_points)

    def generate_double_cables(self):
        """生成2根钢缆（左右分布）"""
        # 两根钢缆的X偏移（左右对称）
        offsets = [
            -self.cable_spacing/2,  # 左钢缆
            self.cable_spacing/2    # 右钢缆
        ]
        
        # 两根钢缆的颜色
        colors = [
            [255, 0, 0],    # 左钢缆红色
            [0, 255, 0]     # 右钢缆绿色
        ]
        
        cables = []
        cable_colors = []
        
        for i, x_off in enumerate(offsets):
            cable = self.generate_single_cable(x_off)
            cables.append(cable)
            
            # 生成对应钢缆的颜色
            color = np.zeros((cable.shape[0], 3), dtype=np.uint8)
            color[:, 0] = colors[i][0]
            color[:, 1] = colors[i][1]
            color[:, 2] = colors[i][2]
            cable_colors.append(color)
            
        return cables, cable_colors

    def generate_cargo_points(self, cables):
        """在两根钢缆末端之间生成货物点云"""
        # 计算两根钢缆末端点的平均位置
        end_points = []
        for cable in cables:
            # 取每根钢缆末端10%点的平均
            end_part = cable[-int(len(cable)*0.1):]
            end_points.append(np.mean(end_part, axis=0))
        
        cargo_center = np.mean(end_points, axis=0)  # 货物中心
        
        # 生成货物范围内的随机点
        x = np.random.uniform(
            cargo_center[0] - 0.6, 
            cargo_center[0] + 0.6, 
            self.cargo_points
        )
        y = np.random.uniform(
            cargo_center[1] - 0.6, 
            cargo_center[1] + 0.6, 
            self.cargo_points
        )
        z = np.random.uniform(
            cargo_center[2], 
            cargo_center[2] + 1.0,  # 货物高度
            self.cargo_points
        )
        
        return np.column_stack((x, y, z))

    def create_colored_pointcloud_msg(self, points, colors):
        """创建带颜色的PointCloud2消息"""
        header = Header()
        header.stamp = self.get_clock().now().to_msg()
        header.frame_id = 'lidar_link'  # 以激光雷达为坐标系原点
        
        num_points = points.shape[0]
        row_step = num_points * COLORED_POINT_STEP
        
        # 构建带颜色的数据缓冲区
        buffer = []
        for i in range(num_points):
            x, y, z = points[i]
            r, g, b = colors[i]
            buffer.append(struct.pack('fffBBB', x, y, z, r, g, b))
        data = b''.join(buffer)
        
        msg = PointCloud2()
        msg.header = header
        msg.height = 1
        msg.width = num_points
        msg.fields = COLORED_FIELDS
        msg.is_bigendian = False
        msg.point_step = COLORED_POINT_STEP
        msg.row_step = row_step
        msg.data = data
        msg.is_dense = False
        return msg

    def timer_callback(self):
        """定时生成并发布带颜色的点云消息"""
        # 生成两根钢缆点云
        cables, cable_colors = self.generate_double_cables()
        
        # 生成货物点云（蓝色）
        cargo_points = self.generate_cargo_points(cables)
        cargo_colors = np.zeros((cargo_points.shape[0], 3), dtype=np.uint8)
        cargo_colors[:, 2] = 255  # 蓝色
        
        # 合并所有点和颜色
        all_points = np.vstack(cables + [cargo_points])
        all_colors = np.vstack(cable_colors + [cargo_colors])
        
        # 发布带颜色的点云
        msg = self.create_colored_pointcloud_msg(all_points, all_colors)
        self.publisher_.publish(msg)
        
        # 计算各部分点数
        left_count, right_count = len(cables[0]), len(cables[1])
        self.get_logger().info(
            f"发布点云: 总点数={all_points.shape[0]} "
            f"(左钢缆: {left_count}, 右钢缆: {right_count}, 货物: {len(cargo_points)})"
        )


def main(args=None):
    rclpy.init(args=args)
    node = CablePointCloudSimulator()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    node.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
